import socket
import select

# Telnet protocol escape codes

class TelnetCodes:
    SE   = 0xF0 # End of subnegotiation
    NOP  = 0xF1 # No operation
    DM   = 0xF2 # Data mark
    BRK  = 0xF3 # Break
    IP   = 0xF4 # Suspend
    AO   = 0xF5 # Abort output
    AYT  = 0xF6 # Are you there
    EC   = 0xF7 # Erase character
    EL   = 0xF8 # Erase line
    GA   = 0xF9 # Go ahead
    SB   = 0xFA # Subnegotiation
    WILL = 0xFB # Will
    WONT = 0xFC # Will
    DO   = 0xFD # Do
    DONT = 0xFE # Don't
    IAC  = 0xFF # Interpret as command

    ECHO = 0x01 # Echo
    SGA  = 0x03 # Supress go ahead
    STAT = 0x05 # Status
    TM   = 0x06 # Timing mark
    TT   = 0x18 # Terminal type
    WS   = 0x1F # Window size
    TS   = 0x20 # Terminal speed
    RFC  = 0x21 # Remote flow control
    LM   = 0x22 # Line mode
    EV   = 0x24 # Environment variables

    bSE   = bytes( [SE  ] ) # End of subnegotiation
    bNOP  = bytes( [NOP ] ) # No operation
    bDM   = bytes( [DM  ] ) # Data mark
    bBRK  = bytes( [BRK ] ) # Break
    bIP   = bytes( [IP  ] ) # Suspend
    bAO   = bytes( [AO  ] ) # Abort output
    bAYT  = bytes( [AYT ] ) # Are you there
    bEC   = bytes( [EC  ] ) # Erase character
    bEL   = bytes( [EL  ] ) # Erase line
    bGA   = bytes( [GA  ] ) # Go ahead
    bSB   = bytes( [SB  ] ) # Subnegotiation
    bWILL = bytes( [WILL] ) # Will
    bWONT = bytes( [WONT] ) # Will
    bDO   = bytes( [DO  ] ) # Do
    bDONT = bytes( [DONT] ) # Don't
    bIAC  = bytes( [IAC ] ) # Interpret as command

    bECHO = bytes( [ECHO] ) # Echo
    bSGA  = bytes( [SGA ] ) # Supress go ahead
    bSTAT = bytes( [STAT] ) # Status
    bTM   = bytes( [TM  ] ) # Timing mark
    bTT   = bytes( [TT  ] ) # Terminal type
    bWS   = bytes( [WS  ] ) # Window size
    bTS   = bytes( [TS  ] ) # Terminal speed
    bRFC  = bytes( [RFC ] ) # Remote flow control
    bLM   = bytes( [LM  ] ) # Line mode
    bEV   = bytes( [EV  ] ) # Environment variables

CodeStrings = {}

CodeStrings[ TelnetCodes.SE   ] = 'SE'   # End of subnegotiation
CodeStrings[ TelnetCodes.NOP  ] = 'NOP'  # No operation
CodeStrings[ TelnetCodes.DM   ] = 'DM'   # Data mark
CodeStrings[ TelnetCodes.BRK  ] = 'BRK'  # Break
CodeStrings[ TelnetCodes.IP   ] = 'IP'   # Suspend
CodeStrings[ TelnetCodes.AO   ] = 'AO'   # Abort output
CodeStrings[ TelnetCodes.AYT  ] = 'AYT'  # Are you there
CodeStrings[ TelnetCodes.EC   ] = 'EC'   # Erase character
CodeStrings[ TelnetCodes.EL   ] = 'EL'   # Erase line
CodeStrings[ TelnetCodes.GA   ] = 'GA'   # Go ahead
CodeStrings[ TelnetCodes.SB   ] = 'SB'   # Subnegotiation
CodeStrings[ TelnetCodes.WILL ] = 'WILL' # Will
CodeStrings[ TelnetCodes.WONT ] = 'WONT' # Will
CodeStrings[ TelnetCodes.DO   ] = 'DO'   # Do
CodeStrings[ TelnetCodes.DONT ] = 'DONT' # Don't
CodeStrings[ TelnetCodes.IAC  ] = 'IAC'  # Interpret as command

CodeStrings[ TelnetCodes.ECHO ] = 'ECHO' # Echo
CodeStrings[ TelnetCodes.SGA  ] = 'SGA'  # Supress go ahead
CodeStrings[ TelnetCodes.STAT ] = 'STAT' # Status
CodeStrings[ TelnetCodes.TM   ] = 'TM'   # Timing mark
CodeStrings[ TelnetCodes.TT   ] = 'TT'   # Terminal type
CodeStrings[ TelnetCodes.WS   ] = 'WS'   # Window size
CodeStrings[ TelnetCodes.TS   ] = 'TS'   # Terminal speed
CodeStrings[ TelnetCodes.RFC  ] = 'RFC'  # Remote flow control
CodeStrings[ TelnetCodes.LM   ] = 'LM'   # Line mode
CodeStrings[ TelnetCodes.EV   ] = 'EV'   # Environment variables

class EchoMode:
    NoEcho   = 0 # Default
    Echo     = 1
    Password = 2 

class TelnetServer:

    def __init__( self, port=8945 ):

        self._port = port

        self._socketData = {}
        self._echoModes = {}

        self._listenSocket = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
        self._listenSocket.setblocking( 0 )
        self._listenSocket.bind( ('', port) )

        self._sockets = [ self._listenSocket ]

    def start( self ):

        """Starts an event loop which constantly listens for new socket
           connections and reads from all connected sockets."""

        self._listenSocket.listen( 1 )

        print( 'Telnet server listening on port ' + str(self._port) )

        while 1: # Main event loop
            try:
                ready_read, ready_write, in_error = select.select( self._sockets, [], [], .1 )

                for sock in ready_read:
                    if sock in self._socketData:

                        try:
                            data = sock.recv( 1024 )
                        except:
                            self._disconnect( sock )
                        else:
                            if data:
                                if data[0] == TelnetCodes.IAC:
                                    self._processIAC( sock, data )

                                else:
                                    if self._echoModes[ sock.fileno() ] == EchoMode.Echo:
                                        sock.send( data )

                                    self._socketData[ sock ] += data.decode()

                                    if self._socketData[ sock ]:
                                        # Don't process data until we've gotten a newline
                                        if self._socketData[ sock ][ -1 ] == '\n':
                                            self.dataEvent( sock.fileno(), self._socketData[sock].rstrip() )
                                            self._socketData[ sock ] = ''

                                    else:
                                        self._disconnect( sock )

                    else:
                        connection, address = sock.accept()
                        print( 'New client: ' + str(address) )
                        self._sockets.append( connection )
                        self._socketData[ connection ] = ''
                        self._echoModes[ connection.fileno() ] = EchoMode.NoEcho

                        self.connectEvent( connection.fileno() )

            except KeyboardInterrupt:
                print( 'Shutting down...' )
                break

            self.loopEvent()

    def _disconnect( self, sock ):
        print( 'Connection closed' )
        self._sockets.remove( sock )
        self.disconnectEvent( sock.fileno() )

    def _processIAC( self, sock, data ):
        if not data or data[0] != TelnetCodes.IAC:
            return
        
        print( "Processing IAC: " + self._translateIAC(data) )

        commands = data.split( TelnetCodes.bIAC )

        response = b''

        for command in commands:
            if not command: continue

            if command == TelnetCodes.bDO + TelnetCodes.bECHO:
                response += TelnetCodes.bIAC + TelnetCodes.bWILL + TelnetCodes.bECHO

                # Don't change the echo mode if we're already echoing so that
                # if we're in password echo mode we don't start echoing password
                # characters
                if self._echoModes[sock.fileno()] == EchoMode.NoEcho:
                    self.setEchoMode( sock.fileno(), EchoMode.Echo )

            elif command == TelnetCodes.bDONT + TelnetCodes.bECHO:
                self.setEchoMode( sock.fileno(), EchoMode.NoEcho )

            # Ignore Process interrupt
            elif command == TelnetCodes.bIP:
                response += TelnetCodes.bIAC + TelnetCodes.bWONT + TelnetCodes.bIP

            # Ignore Timing marks
            elif command == TelnetCodes.bDO + TelnetCodes.bTM:
                response += TelnetCodes.bIAC + TelnetCodes.bWONT + TelnetCodes.bTM

            else:
                print( "Unrecognized command: " + self._translateIAC(command) )

        if response:
            print( "Responding with: " + self._translateIAC(response) )
        else:
            print( "Not responding" )

        sock.send( response )

    def send( self, id, data ):
        """Sends the given data to the client identified by id"""

        sock = self._socketFromId( id )

        try:
            sock.send( data.encode() )
        except AttributeError:
            sock.send( data ) # We received a bytes sequence instead of a string

    def sendLine( self, id, data='' ):
        """Sends the given data to the client identified by id, 
           automatically appending the CR LF characters """

        self.send( id, data + '\r\n' )

    def connectEvent( self, id ):
        """This function is called whenever there is a new connection"""
        pass

    def disconnectEvent( self, id ):
        """This function is called whenever a connection is closed"""
        pass

    def dataEvent( self, id, data ):
        """This function is called whenever data has been read from a socket"""
        pass

    def loopEvent( self ):
        """Called on each pass of the event loop"""
        pass

    def setEchoMode( self, id, mode ):
        if self._echoModes[ id ] == mode:
            return

        if mode == EchoMode.Password:
            self.send( id, TelnetCodes.bIAC  + 
                           TelnetCodes.bWILL + 
                           TelnetCodes.bECHO )

        elif mode == EchoMode.NoEcho:
            self.send( id, TelnetCodes.bIAC  + 
                           TelnetCodes.bWONT + 
                           TelnetCodes.bECHO )

        elif mode == EchoMode.Echo:
            self.send( id, TelnetCodes.bIAC  + 
                           TelnetCodes.bWILL + 
                           TelnetCodes.bECHO )

        self._echoModes[ id ] = mode

    def _socketFromId( self, id ):
        return socket.fromfd( id, socket.AF_INET, socket.SOCK_STREAM )

    def _translateIAC( self, sequence ):
        translation = ''
        for byte in sequence:
            if byte in CodeStrings:
                translation += ( CodeStrings[byte] + ' ' )
            else:
                translation += '?' + str(byte) + ' '
        return translation

